预测可视化

我们可以通过在先前创建的predict.ipynb 中添加以下代码来实现预测可视化:

import matplotlib.pyplot as plt
%matplotlib inline
plt.imshow(result[0].plot()[:,:,::-1])

至此predict.ipynb 应该看起来像以下这样:

from ultralytics import YOLO
model = YOLO("./yolov8n.pt", task="detect")
result = model(source="https://ultralytics.com/images/bus.jpg")

import matplotlib.pyplot as plt
%matplotlib inline
plt.imshow(result[0].plot()[:,:,::-1])

提取预测信息

为方便理解,我们可以先通过在predict.ipynb 加入以下代码来查看预测后的所有信息:

result[0].boxes

运行后应当输出类似以下内容:

ultralytics.engine.results.Boxes object with attributes:

cls: tensor([ 5.,  0.,  0.,  0.,  0., 11.], device='cuda:0')
conf: tensor([0.8733, 0.8657, 0.8527, 0.8253, 0.2612, 0.2554], device='cuda:0')
data: tensor([[2.2871e+01, 2.3126e+02, 8.0499e+02, 7.5684e+02, 8.7335e-01, 5.0000e+00],
        [4.8550e+01, 3.9855e+02, 2.4536e+02, 9.0271e+02, 8.6565e-01, 0.0000e+00],
        [6.6947e+02, 3.9216e+02, 8.0972e+02, 8.7703e+02, 8.5271e-01, 0.0000e+00],
        [2.2152e+02, 4.0580e+02, 3.4497e+02, 8.5754e+02, 8.2526e-01, 0.0000e+00],
        [0.0000e+00, 5.5052e+02, 6.3019e+01, 8.7345e+02, 2.6122e-01, 0.0000e+00],
        [5.8280e-02, 2.5446e+02, 3.2561e+01, 3.2487e+02, 2.5537e-01, 1.1000e+01]], device='cuda:0')
id: None
is_track: False
orig_shape: (1080, 810)
shape: torch.Size([6, 6])
xywh: tensor([[413.9319, 494.0517, 782.1224, 525.5781],
        [146.9529, 650.6292, 196.8052, 504.1582],
        [739.5947, 634.5945, 140.2526, 484.8757],
        [283.2496, 631.6707, 123.4507, 451.7409],
        [ 31.5096, 711.9863,  63.0192, 322.9308],
        [ 16.3095, 289.6657,  32.5024,  70.4110]], device='cuda:0')
xywhn: tensor([[0.5110, 0.4575, 0.9656, 0.4866],
        [0.1814, 0.6024, 0.2430, 0.4668],
        [0.9131, 0.5876, 0.1732, 0.4490],
        [0.3497, 0.5849, 0.1524, 0.4183],
        [0.0389, 0.6592, 0.0778, 0.2990],
...
        [5.9939e-02, 3.6903e-01, 3.0291e-01, 8.3584e-01],
        [8.2650e-01, 3.6311e-01, 9.9966e-01, 8.1207e-01],
        [2.7349e-01, 3.7574e-01, 4.2590e-01, 7.9402e-01],
        [0.0000e+00, 5.0974e-01, 7.7801e-02, 8.0875e-01],
        [7.1951e-05, 2.3561e-01, 4.0198e-02, 3.0081e-01]], device='cuda:0')
Output is truncated. View as a scrollable element or open in a text editor. Adjust cell output settings...

可以点击open in a text editor来查看详细信息

输出中的xywh,xywhn等都是输出的类别,例如只需要获取xywh的信息,则应将result[0].boxes改为以下内容:

result[0].boxes.xywh

此时输出的数据仍然为tensor形式,如果要更改为numpy形式,则应该改为以下代码:

result[0].boxes.xywh.cpu().numpy()